#include "octree.h"
#include <torch/extension.h>
#include <cuda_runtime_api.h>
#include <cooperative_groups.h>
#include <cstdio>
#include <iostream>

namespace cg = cooperative_groups;

#define BLOCK_SIZE 256

namespace OCTREE {

// Number, xyz_min, xyz_max, sub_nodes
std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor>
CreateOctree(const torch::Tensor& points, const torch::Tensor& xyz_min, const torch::Tensor& xyz_max)
{
    // auto xyz_mins = torch::zeros({1, 3}, points.options());
    // auto xyz_maxs = torch::zeros({1, 3}, points.options());
    // auto point_num = torch.zeros({1}, points.options().dtype(torch::kInt32));
    // auto sons = torch::zeros({1, 8}, points.options().dtype(torch::kInt32));
    // auto is_leaf = torch::zeros({1}, points.options().dtype(torch::kBool));
    // xyz_mins.index({0}) = xyz_min;
    // xyz_maxs.index({0}) = xyz_max;
    // point_num.index({0}) = points.size(0);
    // is_leaf.index({0}) = false;

    // std::vector<size_t> now_node_idxs(1, 0);
    // std::vector<float> now_xyz_mins(8*3, 0);
    // std::vector<float> now_xyz_maxs(8*3, 0);
    // std::vector<size_t> now_point_nums(8, 0);
    // std::vector<size_t> now_sons(8, 0);
    // std::vector<bool> now_is_leaf(8, 0);

    // dim3 grid((now_node_idxs.size() + BLOCK_SIZE - 1) / BLOCK_SIZE, 1, 1)
    // dim3 block(BLOCK_SIZE, 1, 1);

    // while(true) {
    //     CreateOctreeCUDA <<<grid, block>>> (
    //         points
    //     )
    // }


    return std::make_tuple(0, torch::zeros({1}, points.options()), torch::zeros({1}, points.options()), torch::zeros({1}, points.options()));
}

__device__ long FindGridWithPoint(
    long idx,
    float3 point,
    const float* xyz_mins,
    const float* xyz_maxs,
    const long* sons,
    const bool* is_leaf_nodes)
{
    float3 xyz_min = {xyz_mins[3*idx], xyz_mins[3*idx+1], xyz_mins[3*idx+2]};
    float3 xyz_max = {xyz_maxs[3*idx], xyz_maxs[3*idx+1], xyz_maxs[3*idx+2]};
    const long* son = sons + idx * 8;

    if(xyz_min.x > point.x || xyz_max.x < point.x ||
       xyz_min.y > point.y || xyz_max.y < point.y ||
       xyz_min.z > point.z || xyz_max.z < point.z) return 0;
    if(is_leaf_nodes[idx]) return idx;

    for(int i = 0; i < 8; i ++) {
        int grid_idx = FindGridWithPoint(son[i], point, xyz_mins, xyz_maxs, sons, is_leaf_nodes);
        if(grid_idx != 0) return grid_idx;
    }

    return -1;
}

__global__ void QueryGridCUDA(int P,
    const float* points,
    const float* xyz_mins,
    const float* xyz_maxs,
    const long* sons,
    const bool* is_leaf_nodes,
    long* grid_idxs)
{
    auto block = cg::this_thread_block();
    auto i = blockIdx.x * blockDim.x + threadIdx.x;

    if(i >= P) return;

    float3 point = {points[3*i], points[3*i+1], points[3*i+2]};

    long grid_idx = FindGridWithPoint(0, point, xyz_mins, xyz_maxs, sons, is_leaf_nodes);
    grid_idxs[i] = grid_idx;
}

torch::Tensor QueryGrid(const torch::Tensor& points, const torch::Tensor& xyz_mins, const torch::Tensor& xyz_maxs, const torch::Tensor& sons, const torch::Tensor& is_leaf_nodes)
{
    auto grid_idxs = torch::zeros({points.size(0)}, points.options().dtype(torch::kLong));

    dim3 grid((points.size(0) + BLOCK_SIZE - 1) / BLOCK_SIZE, 1, 1);
    dim3 block(BLOCK_SIZE, 1, 1);

    QueryGridCUDA <<<grid, block>>> (
        points.size(0),
        points.contiguous().data<float>(),
        xyz_mins.contiguous().data<float>(),
        xyz_maxs.contiguous().data<float>(),
        sons.contiguous().data<long>(),
        is_leaf_nodes.contiguous().data<bool>(),
        grid_idxs.contiguous().data<long>()
    );

    return grid_idxs;
}

} // namespace OCTREE